"""Secure Transport Layer.

Provides encrypted data transmission with replay attack protection
and message integrity checking for BLE communications.
"""

import asyncio
import hashlib
import hmac
import json
import logging
import secrets
import struct
import time
from typing import Dict, List, Optional, Set, Any, Callable
from enum import Enum
from dataclasses import dataclass, field
from collections import defaultdict, deque

from .key_manager import SessionKeyManager, SessionKey

logger = logging.getLogger(__name__)


# Transport Configuration
TRANSPORT_CONFIG = {
    "max_message_size": 8192,         # Maximum message size
    "sequence_window_size": 100,      # Replay protection window
    "timestamp_tolerance": 30,        # Timestamp tolerance (seconds)
    "nonce_cache_size": 1000,        # Nonce cache size
    "cleanup_interval": 60,           # Cleanup interval
    "max_retries": 3,                # Maximum retransmission attempts
    "retry_timeout": 1.0,             # Retry timeout (seconds)
    "compression_threshold": 512      # Compression threshold
}


class MessageType(Enum):
    """Secure message types."""
    DATA = 0x01              # Regular data message
    HANDSHAKE = 0x02         # Key exchange handshake
    HEARTBEAT = 0x03         # Connection heartbeat
    ACK = 0x04              # Acknowledgment
    ERROR = 0x05            # Error message
    CONTROL = 0x06          # Control message


class SecurityLevel(Enum):
    """Message security levels."""
    NONE = 0                # No encryption
    AUTH = 1                # Authentication only
    ENCRYPTED = 2           # Encrypted + authenticated


@dataclass
class MessageHeader:
    """Secure message header."""
    version: int = 1
    msg_type: MessageType = MessageType.DATA
    security_level: SecurityLevel = SecurityLevel.ENCRYPTED
    sequence: int = 0
    timestamp: int = field(default_factory=lambda: int(time.time()))
    session_id: str = ""
    message_id: str = field(default_factory=lambda: secrets.token_urlsafe(8))
    
    def to_bytes(self) -> bytes:
        """Serialize header to bytes."""
        return struct.pack(
            '>BBBIIQ16s8s',
            self.version,
            self.msg_type.value,
            self.security_level.value,
            self.sequence,
            self.timestamp,
            len(self.session_id.encode()),
            self.session_id.encode()[:16].ljust(16, b'\x00'),
            self.message_id.encode()[:8].ljust(8, b'\x00')
        )
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'MessageHeader':
        """Deserialize header from bytes."""
        if len(data) < 43:
            raise ValueError("Invalid header data")
        
        fields = struct.unpack('>BBBIIQ16s8s', data[:43])
        
        header = cls()
        header.version = fields[0]
        header.msg_type = MessageType(fields[1])
        header.security_level = SecurityLevel(fields[2])
        header.sequence = fields[3]
        header.timestamp = fields[4]
        header.session_id = fields[6].rstrip(b'\x00').decode('utf-8')
        header.message_id = fields[7].rstrip(b'\x00').decode('utf-8')
        
        return header


@dataclass
class SecureMessage:
    """Secure message with encryption and authentication."""
    header: MessageHeader
    payload: bytes
    encrypted_data: Optional[bytes] = None
    iv: Optional[bytes] = None
    auth_tag: Optional[bytes] = None
    signature: Optional[bytes] = None
    
    @property
    def is_encrypted(self) -> bool:
        """Check if message is encrypted."""
        return (self.header.security_level == SecurityLevel.ENCRYPTED and
                self.encrypted_data is not None)
    
    @property
    def total_size(self) -> int:
        """Get total message size."""
        size = len(self.header.to_bytes())
        size += len(self.encrypted_data or self.payload)
        if self.iv:
            size += len(self.iv)
        if self.auth_tag:
            size += len(self.auth_tag)
        if self.signature:
            size += len(self.signature)
        return size
    
    def to_wire_format(self) -> bytes:
        """Serialize message for transmission."""
        data = self.header.to_bytes()
        
        if self.is_encrypted:
            # Add encrypted data
            data += struct.pack('>H', len(self.iv)) + self.iv
            data += struct.pack('>H', len(self.auth_tag)) + self.auth_tag
            data += struct.pack('>I', len(self.encrypted_data)) + self.encrypted_data
        else:
            # Add plaintext data
            data += struct.pack('>I', len(self.payload)) + self.payload
        
        if self.signature:
            data += struct.pack('>H', len(self.signature)) + self.signature
        
        return data
    
    @classmethod
    def from_wire_format(cls, data: bytes) -> 'SecureMessage':
        """Deserialize message from transmission format."""
        if len(data) < 43:
            raise ValueError("Invalid message data")
        
        # Parse header
        header = MessageHeader.from_bytes(data[:43])
        offset = 43
        
        message = cls(header=header, payload=b'')
        
        if header.security_level == SecurityLevel.ENCRYPTED:
            # Parse IV
            iv_len = struct.unpack('>H', data[offset:offset+2])[0]
            offset += 2
            message.iv = data[offset:offset+iv_len]
            offset += iv_len
            
            # Parse auth tag
            tag_len = struct.unpack('>H', data[offset:offset+2])[0]
            offset += 2
            message.auth_tag = data[offset:offset+tag_len]
            offset += tag_len
            
            # Parse encrypted data
            data_len = struct.unpack('>I', data[offset:offset+4])[0]
            offset += 4
            message.encrypted_data = data[offset:offset+data_len]
            offset += data_len
        else:
            # Parse plaintext data
            data_len = struct.unpack('>I', data[offset:offset+4])[0]
            offset += 4
            message.payload = data[offset:offset+data_len]
            offset += data_len
        
        # Parse signature if present
        if offset < len(data):
            sig_len = struct.unpack('>H', data[offset:offset+2])[0]
            offset += 2
            message.signature = data[offset:offset+sig_len]
        
        return message


class ReplayProtection:
    """Replay attack protection."""
    
    def __init__(self):
        self.sequence_windows: Dict[str, Set[int]] = defaultdict(set)
        self.nonce_cache: Set[str] = set()
        self.last_sequences: Dict[str, int] = defaultdict(int)
        self.last_cleanup = time.time()
    
    def validate_sequence(self, session_id: str, sequence: int) -> bool:
        """Validate message sequence number."""
        window = self.sequence_windows[session_id]
        window_size = TRANSPORT_CONFIG["sequence_window_size"]
        last_seq = self.last_sequences[session_id]
        
        # Check if sequence is within acceptable window
        if sequence <= last_seq - window_size:
            return False  # Too old
        
        if sequence in window:
            return False  # Duplicate
        
        # Add to window
        window.add(sequence)
        
        # Update last sequence if newer
        if sequence > last_seq:
            self.last_sequences[session_id] = sequence
        
        # Trim window to size
        if len(window) > window_size:
            min_seq = max(0, sequence - window_size)
            window = {s for s in window if s > min_seq}
            self.sequence_windows[session_id] = window
        
        return True
    
    def validate_timestamp(self, timestamp: int) -> bool:
        """Validate message timestamp."""
        current_time = int(time.time())
        tolerance = TRANSPORT_CONFIG["timestamp_tolerance"]
        
        return abs(current_time - timestamp) <= tolerance
    
    def validate_nonce(self, nonce: str) -> bool:
        """Validate message nonce."""
        if nonce in self.nonce_cache:
            return False
        
        self.nonce_cache.add(nonce)
        
        # Limit cache size
        cache_size = TRANSPORT_CONFIG["nonce_cache_size"]
        if len(self.nonce_cache) > cache_size:
            # Remove oldest entries (simplified LRU)
            excess = len(self.nonce_cache) - cache_size + 100
            for _ in range(excess):
                self.nonce_cache.pop()
        
        return True
    
    def cleanup_expired(self):
        """Clean up expired data."""
        current_time = time.time()
        if current_time - self.last_cleanup < TRANSPORT_CONFIG["cleanup_interval"]:
            return
        
        # Clean up old sequences
        cutoff_time = current_time - TRANSPORT_CONFIG["timestamp_tolerance"] * 2
        
        # This is a simplified cleanup - in production you'd want more sophisticated logic
        self.last_cleanup = current_time


@dataclass
class TransmissionRecord:
    """Record of transmitted message for retry logic."""
    message_id: str
    session_id: str
    message: SecureMessage
    sent_time: float
    retry_count: int = 0
    ack_received: bool = False
    
    @property
    def needs_retry(self) -> bool:
        """Check if message needs retry."""
        if self.ack_received:
            return False
        
        timeout = TRANSPORT_CONFIG["retry_timeout"] * (2 ** self.retry_count)
        return (time.time() - self.sent_time) > timeout
    
    @property
    def max_retries_reached(self) -> bool:
        """Check if maximum retries reached."""
        return self.retry_count >= TRANSPORT_CONFIG["max_retries"]


class SecureTransport:
    """Secure transport layer for BLE communications."""
    
    def __init__(self, key_manager: SessionKeyManager):
        self.key_manager = key_manager
        self.replay_protection = ReplayProtection()
        
        # Message tracking
        self.sequence_counters: Dict[str, int] = defaultdict(int)
        self.pending_transmissions: Dict[str, TransmissionRecord] = {}
        self.message_callbacks: Dict[MessageType, Callable] = {}
        
        # Background tasks
        self.retry_task: Optional[asyncio.Task] = None
        self.cleanup_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        logger.info("SecureTransport initialized")
    
    async def start(self):
        """Start secure transport."""
        if self.is_running:
            return
        
        self.is_running = True
        self.retry_task = asyncio.create_task(self._retry_loop())
        self.cleanup_task = asyncio.create_task(self._cleanup_loop())
        
        logger.info("SecureTransport started")
    
    async def stop(self):
        """Stop secure transport."""
        self.is_running = False
        
        if self.retry_task:
            self.retry_task.cancel()
            try:
                await self.retry_task
            except asyncio.CancelledError:
                pass
        
        if self.cleanup_task:
            self.cleanup_task.cancel()
            try:
                await self.cleanup_task
            except asyncio.CancelledError:
                pass
        
        logger.info("SecureTransport stopped")
    
    def register_message_handler(self, msg_type: MessageType, 
                                callback: Callable[[SecureMessage], None]):
        """Register message type handler."""
        self.message_callbacks[msg_type] = callback
    
    def _get_next_sequence(self, session_id: str) -> int:
        """Get next sequence number for session."""
        seq = self.sequence_counters[session_id]
        self.sequence_counters[session_id] = (seq + 1) % 0xFFFFFFFF
        return seq
    
    async def encrypt_message(self, session_id: str, 
                            message: SecureMessage) -> bool:
        """Encrypt message payload."""
        if message.header.security_level != SecurityLevel.ENCRYPTED:
            return True
        
        # Encrypt with session key
        encrypted = self.key_manager.encrypt_data(session_id, message.payload)
        if not encrypted:
            logger.error(f"Encryption failed for session {session_id}")
            return False
        
        # Store encrypted data
        message.iv = bytes.fromhex(encrypted["iv"])
        message.auth_tag = bytes.fromhex(encrypted["tag"])
        message.encrypted_data = bytes.fromhex(encrypted["ciphertext"])
        
        return True
    
    async def decrypt_message(self, session_id: str, 
                            message: SecureMessage) -> bool:
        """Decrypt message payload."""
        if message.header.security_level != SecurityLevel.ENCRYPTED:
            return True
        
        if not message.is_encrypted:
            logger.error("Message marked as encrypted but no encrypted data")
            return False
        
        # Prepare encrypted data
        encrypted_data = {
            "algorithm": "AES-128-GCM",  # From key manager
            "iv": message.iv.hex(),
            "ciphertext": message.encrypted_data.hex(),
            "tag": message.auth_tag.hex()
        }
        
        # Decrypt with session key
        plaintext = self.key_manager.decrypt_data(session_id, encrypted_data)
        if plaintext is None:
            logger.error(f"Decryption failed for session {session_id}")
            return False
        
        message.payload = plaintext
        return True
    
    def validate_message(self, message: SecureMessage) -> bool:
        """Validate message security properties."""
        header = message.header
        
        # Validate timestamp
        if not self.replay_protection.validate_timestamp(header.timestamp):
            logger.warning(f"Message timestamp validation failed: {header.message_id}")
            return False
        
        # Validate sequence
        if not self.replay_protection.validate_sequence(
            header.session_id, header.sequence
        ):
            logger.warning(f"Message sequence validation failed: {header.message_id}")
            return False
        
        # Validate nonce (message ID)
        if not self.replay_protection.validate_nonce(header.message_id):
            logger.warning(f"Message nonce validation failed: {header.message_id}")
            return False
        
        return True
    
    async def send_message(self, session_id: str, payload: bytes,
                          msg_type: MessageType = MessageType.DATA,
                          security_level: SecurityLevel = SecurityLevel.ENCRYPTED,
                          require_ack: bool = True) -> Optional[str]:
        """Send secure message."""
        try:
            # Create message header
            header = MessageHeader(
                msg_type=msg_type,
                security_level=security_level,
                sequence=self._get_next_sequence(session_id),
                session_id=session_id
            )
            
            # Create message
            message = SecureMessage(
                header=header,
                payload=payload
            )
            
            # Check message size
            if len(payload) > TRANSPORT_CONFIG["max_message_size"]:
                logger.error(f"Message too large: {len(payload)} bytes")
                return None
            
            # Encrypt message
            if not await self.encrypt_message(session_id, message):
                return None
            
            # Serialize message
            wire_data = message.to_wire_format()
            
            # Track for retry if ACK required
            if require_ack and msg_type != MessageType.ACK:
                record = TransmissionRecord(
                    message_id=header.message_id,
                    session_id=session_id,
                    message=message,
                    sent_time=time.time()
                )
                self.pending_transmissions[header.message_id] = record
            
            # TODO: Actually send via BLE - this would integrate with BLE stack
            logger.debug(f"Sending {len(wire_data)} bytes for message {header.message_id}")
            
            return header.message_id
            
        except Exception as e:
            logger.error(f"Send message error: {e}")
            return None
    
    async def receive_message(self, wire_data: bytes) -> Optional[SecureMessage]:
        """Receive and validate secure message."""
        try:
            # Deserialize message
            message = SecureMessage.from_wire_format(wire_data)
            
            # Validate message
            if not self.validate_message(message):
                return None
            
            # Decrypt message
            if not await self.decrypt_message(message.header.session_id, message):
                return None
            
            # Handle ACK messages
            if message.header.msg_type == MessageType.ACK:
                self._handle_ack(message)
                return message
            
            # Send ACK if required
            if message.header.msg_type != MessageType.HEARTBEAT:
                await self._send_ack(message)
            
            # Call message handler
            handler = self.message_callbacks.get(message.header.msg_type)
            if handler:
                try:
                    handler(message)
                except Exception as e:
                    logger.error(f"Message handler error: {e}")
            
            logger.debug(f"Received message {message.header.message_id} "
                        f"from session {message.header.session_id}")
            
            return message
            
        except Exception as e:
            logger.error(f"Receive message error: {e}")
            return None
    
    def _handle_ack(self, ack_message: SecureMessage):
        """Handle ACK message."""
        # ACK payload should contain the original message ID
        if len(ack_message.payload) >= 8:
            orig_message_id = ack_message.payload.decode('utf-8', errors='ignore')
            record = self.pending_transmissions.get(orig_message_id)
            if record:
                record.ack_received = True
                logger.debug(f"ACK received for message {orig_message_id}")
    
    async def _send_ack(self, original_message: SecureMessage) -> bool:
        """Send ACK for received message."""
        try:
            ack_payload = original_message.header.message_id.encode('utf-8')
            
            await self.send_message(
                original_message.header.session_id,
                ack_payload,
                MessageType.ACK,
                SecurityLevel.AUTH,
                require_ack=False
            )
            
            return True
            
        except Exception as e:
            logger.error(f"Send ACK error: {e}")
            return False
    
    async def send_heartbeat(self, session_id: str) -> bool:
        """Send heartbeat message."""
        heartbeat_data = json.dumps({
            "timestamp": int(time.time()),
            "session_id": session_id
        }).encode('utf-8')
        
        message_id = await self.send_message(
            session_id, 
            heartbeat_data,
            MessageType.HEARTBEAT,
            SecurityLevel.AUTH,
            require_ack=False
        )
        
        return message_id is not None
    
    async def _retry_loop(self):
        """Background retry loop for failed transmissions."""
        while self.is_running:
            try:
                await asyncio.sleep(1.0)
                
                # Check pending transmissions
                retry_messages = []
                expired_messages = []
                
                for msg_id, record in self.pending_transmissions.items():
                    if record.ack_received:
                        continue
                    
                    if record.max_retries_reached:
                        expired_messages.append(msg_id)
                    elif record.needs_retry:
                        retry_messages.append(msg_id)
                
                # Remove expired messages
                for msg_id in expired_messages:
                    del self.pending_transmissions[msg_id]
                    logger.warning(f"Message transmission failed: {msg_id}")
                
                # Retry messages
                for msg_id in retry_messages:
                    record = self.pending_transmissions[msg_id]
                    record.retry_count += 1
                    record.sent_time = time.time()
                    
                    # TODO: Resend message via BLE
                    logger.debug(f"Retrying message {msg_id} (attempt {record.retry_count})")
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Retry loop error: {e}")
    
    async def _cleanup_loop(self):
        """Background cleanup loop."""
        while self.is_running:
            try:
                await asyncio.sleep(TRANSPORT_CONFIG["cleanup_interval"])
                
                # Clean up completed transmissions
                completed = [msg_id for msg_id, record in self.pending_transmissions.items()
                           if record.ack_received]
                for msg_id in completed:
                    del self.pending_transmissions[msg_id]
                
                # Clean up replay protection
                self.replay_protection.cleanup_expired()
                
                if completed:
                    logger.debug(f"Cleaned up {len(completed)} completed transmissions")
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Cleanup loop error: {e}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get transport statistics."""
        pending_count = len(self.pending_transmissions)
        retry_count = sum(1 for r in self.pending_transmissions.values() 
                         if r.retry_count > 0)
        
        return {
            "secure_transport": {
                "is_running": self.is_running,
                "pending_transmissions": pending_count,
                "retrying_messages": retry_count,
                "active_sessions": len(self.sequence_counters)
            },
            "replay_protection": {
                "tracked_sessions": len(self.replay_protection.sequence_windows),
                "nonce_cache_size": len(self.replay_protection.nonce_cache)
            }
        }
    
    async def cleanup(self):
        """Clean up secure transport."""
        await self.stop()
        self.pending_transmissions.clear()
        self.sequence_counters.clear()
        logger.info("SecureTransport cleanup completed")


# Example usage and testing
if __name__ == "__main__":
    import asyncio
    from .key_manager import SessionKeyManager
    
    async def test_secure_transport():
        """Test secure transport functionality."""
        # Create key manager and transport
        key_manager = SessionKeyManager()
        transport = SecureTransport(key_manager)
        
        # Start services
        await key_manager.start()
        await transport.start()
        
        try:
            session_id = "test_session_001"
            
            # Generate session key
            session_key = await key_manager.generate_session_key(session_id, "test123")
            print(f"Session key generated: {session_key.key_id}")
            
            # Test message transmission
            test_data = b"Hello, secure world!"
            print(f"Sending message: {test_data}")
            
            message_id = await transport.send_message(
                session_id, test_data, MessageType.DATA
            )
            
            if message_id:
                print(f"Message sent with ID: {message_id}")
                
                # Simulate message creation for testing
                header = MessageHeader(
                    msg_type=MessageType.DATA,
                    sequence=0,
                    session_id=session_id,
                    message_id=message_id
                )
                
                test_message = SecureMessage(header=header, payload=test_data)
                
                # Test encryption
                if await transport.encrypt_message(session_id, test_message):
                    print("Message encrypted successfully")
                    print(f"IV length: {len(test_message.iv)} bytes")
                    print(f"Encrypted data length: {len(test_message.encrypted_data)} bytes")
                    
                    # Test decryption
                    if await transport.decrypt_message(session_id, test_message):
                        print("Message decrypted successfully")
                        print(f"Decrypted: {test_message.payload}")
                        print(f"Match: {test_message.payload == test_data}")
                    else:
                        print("Decryption failed")
                else:
                    print("Encryption failed")
            else:
                print("Message send failed")
            
            # Test heartbeat
            print("\nTesting heartbeat...")
            heartbeat_sent = await transport.send_heartbeat(session_id)
            print(f"Heartbeat sent: {heartbeat_sent}")
            
            # Print statistics
            stats = transport.get_statistics()
            print(f"\nTransport statistics: {json.dumps(stats, indent=2)}")
            
        finally:
            # Cleanup
            await transport.cleanup()
            await key_manager.cleanup()
    
    # Run test
    asyncio.run(test_secure_transport())